import types, hashlib
from Crypto.Cipher import AES

class Vault:
    """A secure vault that can be used to store information secured
    with a password.
    
    You can create a Vault by first importing it:
        from Vault import Vault
    For this to work the Vault.py file must be in the same folder
    as your cracker. If instead you have a script running inside
    your own folder then add this to the top of your python code:
        import sys
        sys.path.insert(0, "..")
        from Vault import Vault
    I hope that makes sense.
    
    To use a Vault you can simply create one with a password
    and some data like this:
        v = Vault("password123", "Super secret information")
    You can then open the vault to obtain your data like this:
        data = v.open("password123")
    If you get the password wrong then an exception will be raised.
    
    If you have created a Vault and want to send it to a friend securely
    then you can save your encyrpted vault to a file like so:
        v = Vault("password123", "Super secret information")
        v.save("path/to/file.vault")
    If the path you provide points to an existing file it will be
    overwritten
    
    If you been given a Vault that you would like to open then you can
    load it from a file like so:
        v = Vault().load("path/to/file.vault")
    You'll still need the password to get the data out however.
    
    Vaults require the use of the PyCrypto module found here:
        https://www.dlitz.net/software/pycrypto/
    """
    
    magicphrase = "Why is a raven like a writing desk?"
    encodedphrase = ""
    contents = ""
    
    def __init__(self, password="", data=""):
        if type(password) != types.StringType:
            raise TypeError("A Vault's key must be a string")
        self.encodedphrase = self.encode(password, self.magicphrase)
        self.contents = self.encode(password, data)
        
    def open(self, password=""):
        key = hashlib.sha256(password).digest()
        phrase = self.decode(password, self.encodedphrase)
        if phrase == self.magicphrase:
            return self.decode(password, self.contents)
        else:
            raise Exception("Incorrect password")
            
    def encode(self, password, plaintext):
        size = len(plaintext)
        plaintext = plaintext.zfill(((size+15)/16)*16)
        key = hashlib.sha256(password).digest()
        mode = AES.MODE_CBC
        encryptor = AES.new(key, mode)
        return str(size) + ":" + encryptor.encrypt(plaintext)
        
    def decode(self, password, encodedtext):
        size, colon, encodedtext = encodedtext.partition(":")
        key = hashlib.sha256(password).digest()
        mode = AES.MODE_CBC
        decryptor = AES.new(key, mode)
        plaintext = decryptor.decrypt(encodedtext)
        return plaintext[len(encodedtext) - int(size):]
        
    def save(self, path):
        file = open(path, "wb")
        file.write(self.encodedphrase)
        file.write("\n")
        file.write(self.contents)
        file.close()
        
    def load(self, path):
        file = open(path, "rb")
        self.encodedphrase = file.readline()[0:-1]
        self.contents = file.read()
        file.close()
        return self

def test():
    # Tests

    # Create vault with no password
    try:
        nopassword = Vault(data="no password")
    except:
        print "Test failure: Should be able to create a vault with no password"
    
    # Create vault with no data
    try:
        nodata = Vault(password="123")
    except:
        print "Test failure: Should be able to create a vault with no data"
    
    # Create vault with non-string password
    try:
        v = Vault(123, "Incorrect")
    except(TypeError):
        pass
    else:
        print "Test failure: Should not be able to create a vault with a non-string key"
    
    # Create vault
    try:
        vault = Vault("123", "Super secret data")
    except:
        print "Test failure: Should be able to create a vault"
    
    # Open a vault with no password
    try:
        nopassword = Vault(data="no password")
        nopassword.open()
    except:
        print "Test failure: Should be able to open a vault with no password"
          
    # Open a vault with no data
    try:
        nodata = Vault(password="123")
        nodata.open("123")
    except:
        print "Test failure: Should be able to open a vault with no data"
     
    # Get correct details
    vault = Vault("123", "Super secret data")
    if "Super secret data" != vault.open("123"):
        print "Test failure: Should get data when key is correct"
                   
    # Deny details
    try:
        vault.open("321")
    except:
        pass
    else:
        print "Test failure: Should not get data when key is incorrect"
        
    # Save Vault to file
    vault.save("test.vault")
    
    # Load from save
    w = Vault().load("test.vault")
    if "Super secret data" != w.open("123"):
        print "Test failure: Should get data from loaded Vault with correct key"
            
if __name__ == "__main__":
    test()
